# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Video models."""

import math
import mindspore.nn as nn
import mindvision.video.utils.x3d_weight_init_helper as init_helper
from mindvision.video.utils.x3d_utils import round_width
from mindvision.video.models.head import x3d_head as head_helper
from mindvision.video.models.blocks import x3d_resnet as resnet_helper
from mindvision.video.models.blocks import x3d_stem as stem_helper
from mindvision.engine.class_factory import ClassFactory, ModuleType

# Number of blocks for different stages given the model depth.
_MODEL_STAGE_DEPTH = {50: (3, 4, 6, 3), 101: (3, 4, 23, 3)}

# Basis of temporal kernel sizes for each of the stage.
_TEMPORAL_KERNEL_BASIS = {
    "x3d": [
        [[5]],  # conv1 temporal kernels.
        [[3]],  # res2 temporal kernels.
        [[3]],  # res3 temporal kernels.
        [[3]],  # res4 temporal kernels.
        [[3]],  # res5 temporal kernels.
    ],
}

_POOL1 = {
    "x3d": [[1, 1, 1]],
}


@ClassFactory.register(ModuleType.BACKBONE)
class X3D(nn.Cell):
    """
    X3D model builder. It builds a X3D network backbone, which is a ResNet.

    Christoph Feichtenhofer.
    "X3D: Expanding Architectures for Efficient Video Recognition."
    https://arxiv.org/abs/2004.04730

    Args:
        cfg (CfgNode): model building configs, details are in the
            comments of the config file.
    """

    def __init__(self):
        super(X3D, self).__init__()
        # self.norm_module = get_norm(cfg)
        self.norm_module = nn.BatchNorm3d
        self.enable_detection = False
        self.num_pathways = 1

        exp_stage = 2.0
        self.dim_c1 = 12

        self.dim_res2 = self.dim_c1
        self.dim_res3 = round_width(self.dim_res2, exp_stage, divisor=8)
        self.dim_res4 = round_width(self.dim_res3, exp_stage, divisor=8)
        self.dim_res5 = round_width(self.dim_res4, exp_stage, divisor=8)

        self.block_basis = [
            # blocks, c, stride
            [1, self.dim_res2, 2],
            [2, self.dim_res3, 2],
            [5, self.dim_res4, 2],
            [3, self.dim_res5, 2],
        ]
        # self._construct_network(cfg)
        self.seq = nn.SequentialCell()
        self._construct_network()
        # init_helper
        init_helper.init_weights(
            self, None, True
            # self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
        )

    def _round_repeats(self, repeats, multiplier):
        """Round number of layers based on depth multiplier."""
        multiplier = multiplier
        if not multiplier:
            return repeats
        return int(math.ceil(multiplier * repeats))

    # def _construct_network(self, cfg):
    def _construct_network(self):
        """
        Builds a single pathway X3D model.

        Args:
            cfg (CfgNode): model building configs, details are in the
                comments of the config file.
        """
        assert 'x3d' in _POOL1.keys()

        # num_groups = cfg.RESNET.NUM_GROUPS
        # width_per_group = cfg.RESNET.WIDTH_PER_GROUP
        # dim_inner = num_groups * width_per_group

        # w_mul = cfg.X3D.WIDTH_FACTOR
        # d_mul = cfg.X3D.DEPTH_FACTOR
        w_mul = 2.0
        d_mul = 2.2
        dim_res1 = round_width(self.dim_c1, w_mul)

        temp_kernel = _TEMPORAL_KERNEL_BASIS['x3d']

        s1 = stem_helper.VideoModelStem(
            # dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
            dim_in=[3],
            dim_out=[dim_res1],
            kernel=[temp_kernel[0][0] + [3, 3]],
            stride=[[1, 2, 2]],
            padding=[[temp_kernel[0][0][0] // 2, 1, 1]],
            norm_module=self.norm_module
        )
        self.seq.append(s1)

        dim_in = dim_res1

        for stage, block in enumerate(self.block_basis):
            dim_out = round_width(block[1], w_mul)
            # dim_inner = int(cfg.X3D.BOTTLENECK_FACTOR * dim_out)
            dim_inner = int(2.25 * dim_out)
            n_rep = self._round_repeats(block[0], d_mul)

            s = resnet_helper.ResStage(
                dim_in=[dim_in],
                dim_out=[dim_out],
                dim_inner=[dim_inner],
                temp_kernel_sizes=temp_kernel[1],
                stride=[block[2]],
                num_blocks=[n_rep],
                # num_groups=[dim_inner],
                num_block_temp_kernel=[n_rep],
                stride_1x1=False,
                norm_module=self.norm_module,
                dilation=[1, 1],
                drop_connect_rate=0.5 * (stage + 2)
                / (len(self.block_basis) + 1),
            )
            dim_in = dim_out
            self.seq.append(s)

        # spat_sz = int(math.ceil(cfg.DATA.TRAIN_CROP_SIZE / 32.0))
        spat_sz = int(math.ceil(224 / 32.0))
        head = head_helper.X3DHead(
            dim_in=dim_out,
            dim_inner=dim_inner,
            # dim_out=cfg.X3D.DIM_C5,
            dim_out=2048,
            # num_classes=cfg.MODEL.NUM_CLASSES,
            num_classes=400,
            # pool_size=[cfg.DATA.NUM_FRAMES, spat_sz, spat_sz],
            pool_size=[16, spat_sz, spat_sz],
            # dropout_rate=cfg.MODEL.DROPOUT_RATE,
            dropout_rate=0.5,
        )
        self.seq.append(head)

    def construct(self, x):
        x = self.seq(x)
        return x
